#include "udp.h"
#include "ip.h"
#include "icmp.h"

/**
 * @brief udp处理程序表
 * 
 */
map_t udp_table;

/**
 * @brief udp伪校验和计算
 * 
 * @param buf 要计算的包
 * @param src_ip 源ip地址
 * @param dst_ip 目的ip地址
 * @return uint16_t 伪校验和
 */
static uint16_t udp_checksum(buf_t *buf, uint8_t *src_ip, uint8_t *dst_ip)
{
    int udp_peso_hdr_size = sizeof(udp_peso_hdr_t);
    uint16_t origin_buf_size = ((udp_hdr_t *)(buf->data))->total_len16;

    // Step1: 增加UDP伪头部
    buf_add_header(buf, udp_peso_hdr_size);
    // Step2: 将被伪头部覆盖的IP头部的部分拷贝出来
    uint8_t ip_hdr_part[udp_peso_hdr_size];
    memcpy(ip_hdr_part, buf->data, udp_peso_hdr_size);
    // Step3: 填写伪头部的字段,协议号必是UDP
    udp_peso_hdr_t *hdr = (udp_peso_hdr_t *)(buf->data);
    memcpy(hdr->src_ip, src_ip, NET_IP_LEN);
    memcpy(hdr->dst_ip, dst_ip, NET_IP_LEN);
    hdr->placeholder = 0;
    hdr->protocol = NET_PROTOCOL_UDP;
    hdr->total_len16 = origin_buf_size;
    // Step4: 检查是否对齐2Byte 计算UDP校验和
    int isOdd = (buf->len % 2 == 1);
    if(isOdd){
        buf_add_padding(buf, 1);
    }
    uint16_t checksum16_temp = checksum16((uint16_t *)(buf->data), buf->len);
    // Step5: 将Step2暂存的IP头部拷贝回来
    memcpy(buf->data, ip_hdr_part, udp_peso_hdr_size);
    // Step6: 去除UDP伪头部和尾部填充(若有)
    buf_remove_header(buf, udp_peso_hdr_size);
    if(isOdd){
        buf_remove_padding(buf, 1);
    }
    // Step7: 返回计算出来的校验和值
    return checksum16_temp;
}

/**
 * @brief 处理一个收到的udp数据包
 * 
 * @param buf 要处理的包
 * @param src_ip 源ip地址
 */
void udp_in(buf_t *buf, uint8_t *src_ip)
{
    // Step1: 包检查
    udp_hdr_t *hdr = (udp_hdr_t *)(buf->data); // 这里...
    if(buf->len < sizeof(udp_hdr_t) || buf->len < swap16(hdr->total_len16)){
        printf("\nudp_in: UDP packet from ip(%s) check size error\n", iptos(src_ip));
        printf("buf_len = %lld, total_len from UDP header = %d\n", buf->len, swap16(hdr->total_len16));
        return;
    }
    // Step2: 重新计算并比较校验和
    uint16_t checksum16_temp = hdr->checksum16;
    hdr->checksum16 = 0;
    uint16_t temp = udp_checksum(buf, src_ip, net_if_ip);
    if(temp != checksum16_temp){
        printf("\nudp_in: UDP packet from ip(%s) check checksum error\n", iptos(src_ip));
        printf("after swap16: hdr->checksum16=0x%x, cal_checksum=0x%x\n", swap16(checksum16_temp), temp);
        return;
    }
    hdr->checksum16 = checksum16_temp;

    // Step3: 查询目的端口号是否有回调函数
    uint16_t dst_port16 = swap16(hdr->dst_port16);
    udp_handler_t *handler = map_get(&udp_table, &dst_port16);
    // Step4:若找到,则去掉UDP,处理
    if(handler){
        printf("\nudp_in: exists dst_port(%d) handler\n", hdr->dst_port16);
        uint16_t src_port16 = swap16(hdr->src_port16);
        buf_remove_header(buf, sizeof(udp_hdr_t));
        (*handler)(buf->data, buf->len, src_ip, src_port16);
        return;
    }
    // Step5: 若没找到,buf_add_header增加IPv4头部,icmp_unreachable发送端口不可达ICMP差错报文
    printf("\nudp_in: not exists dst_port(%d) handler\n", hdr->dst_port16);
    buf_add_header(buf, sizeof(ip_hdr_t));
    printf("\nudp_in: send icmp_unreachable response\n");
    icmp_unreachable(buf, src_ip, ICMP_CODE_PORT_UNREACH);
    return;
}

/**
 * @brief 处理一个要发送的数据包
 * 
 * @param buf 要处理的包
 * @param src_port 源端口号
 * @param dst_ip 目的ip地址
 * @param dst_port 目的端口号
 */
void udp_out(buf_t *buf, uint16_t src_port, uint8_t *dst_ip, uint16_t dst_port)
{
    // Step1: 添加UDP报头
    buf_add_header(buf, sizeof(udp_hdr_t));
    // Step2: 填充UDP首部字段
    udp_hdr_t *hdr = (udp_hdr_t *)(buf->data);
    hdr->src_port16 = swap16(src_port);
    hdr->dst_port16 = swap16(dst_port);
    hdr->total_len16 = swap16(buf->len);
    // Step3: 计算校验和
    hdr->checksum16 = 0;
    hdr->checksum16 = udp_checksum(buf, net_if_ip, dst_ip); // debug发现不用swap16,太迷了,而上面的total_len16又需要...
    // Step4: ip_out
    printf("\nudp_out: send a udp packet\n");
    ip_out(buf, dst_ip, NET_PROTOCOL_UDP);
}

/**
 * @brief 初始化udp协议
 * 
 */
void udp_init()
{
    map_init(&udp_table, sizeof(uint16_t), sizeof(udp_handler_t), 0, 0, NULL);
    net_add_protocol(NET_PROTOCOL_UDP, udp_in);
}

/**
 * @brief 打开一个udp端口并注册处理程序
 * 
 * @param port 端口号
 * @param handler 处理程序
 * @return int 成功为0，失败为-1
 */
int udp_open(uint16_t port, udp_handler_t handler)
{
    printf("\nudp open: port = %d\n", port);
    return map_set(&udp_table, &port, &handler);
}

/**
 * @brief 关闭一个udp端口
 * 
 * @param port 端口号
 */
void udp_close(uint16_t port)
{
    map_delete(&udp_table, &port);
}

/**
 * @brief 发送一个udp包
 * 
 * @param data 要发送的数据
 * @param len 数据长度
 * @param src_port 源端口号
 * @param dst_ip 目的ip地址
 * @param dst_port 目的端口号
 */
void udp_send(uint8_t *data, uint16_t len, uint16_t src_port, uint8_t *dst_ip, uint16_t dst_port)
{
    buf_init(&txbuf, len);
    memcpy(txbuf.data, data, len);
    udp_out(&txbuf, src_port, dst_ip, dst_port);
}